#include <protomol/integrator/normal/NormalModeLangMF.h>
#include <protomol/base/Report.h>
#include <protomol/type/ScalarStructure.h>
#include <protomol/type/Vector3DBlock.h>
#include <protomol/force/ForceGroup.h>
#include <protomol/topology/GenericTopology.h>
#include <protomol/topology/TopologyUtilities.h>
#include <protomol/base/PMConstants.h>
#include <protomol/ProtoMolApp.h>

#include <protomol/integrator/normal/ModifierForceProjection.h>


using namespace std;

using namespace ProtoMol::Report;

using std::string;
using std::vector;

namespace ProtoMol {
  //__________________________________________________ NormalModeCompLang

  const string NormalModeLangMF::keyword( "NormalModeLangMF" );

  NormalModeLangMF::NormalModeLangMF() : STSIntegrator(), NormalModeUtilities()
  {
  }

//Modify the constructor 
  NormalModeLangMF::NormalModeLangMF(Real timestep, int firstMode, int nummode, Real gamma, int seed, Real temperature, bool gencn,
                     ForceGroup *overloadedForces) 
    : STSIntegrator(timestep, overloadedForces), 
// FIXME: check initialization, make sure we aren't doing anything funky
// example: what is 'nummode' used for?
        NormalModeUtilities( firstMode, nummode, gamma, seed, temperature), genCompNoise(gencn)    
  {
  }


  NormalModeLangMF::~NormalModeLangMF() 
  {  
  }

//add to initialize
  void NormalModeLangMF::initialize(ProtoMolApp *app){
    STSIntegrator::initialize(app);
    initializeForces();
    //check valid eigenvectors
    //NM initialization if OK
    int nm_flags = NO_NM_FLAGS;
    if(genCompNoise) nm_flags |= GEN_COMP_NOISE;
    NormalModeUtilities::initialize((int)app->positions.size(), app, myForces, nm_flags); //last int for no complimentary forces or gen noise: GEN_COMP_NOISE
    slowForces.resize(_N);
    slowForcesAvg.resize(_N);
    firstFastMode = nummode + (firstMode - 1);
    numFastMode = _3N - firstFastMode + 1;
    numSlowMode = nummode;
    firstSlowMode = firstMode - 1;

  }

  // FIXME: validate this function
  //Find forces acting inside subspace
  Vector3DBlock* NormalModeLangMF::fastspaceForce(Vector3DBlock * force, Vector3DBlock * iPforce){
    //
    if (force != iPforce)
      iPforce->intoAssign(*force);
    //vector3DBlockTOvect(force, tmpFX);        //put positions-x_0 into linear array
    //calculate M^{1/2}QQ^TM^{-1/2}f using BLAS
    //f'=M^{-1/2}*f
    for( int i=0; i < _3N; i++) {
            iPforce->c[i] *= invSqrtMass[i/3];
    }
    //c=Q^T*M^{-1/2}*f
    char transA = 'T';    
// _rfM is the last rows of interest in Q
// _rfM-(firstFastMode-1) is the number of rows we need to define our subspace
// &((*Q)[_3N*(firstFastMode-1)]) (argument for degemv) lists the first row of interest

                                              // Transpose, LAPACK checks only first character N/V
    int m = _3N; int n = numFastModes; int incxy = 1;     //sizes
    double alpha = 1.0; double beta = 0.0;
    Lapack::dgemv(&transA, &m, &n, &alpha, &((*Q)[_3N*(firstFastMode-1)]), &m, iPforce->c, &incxy, &beta, tmpC, &incxy);

    //f''=Qc
    char transB = 'N'; /* LAPACK checks only first character N/V */
    alpha = 1.0;        beta = 0.0;

    Lapack::dgemv(&transB, &m, &n, &alpha, &((*Q)[_3N*(firstFastMode-1)]), &m, tmpC, &incxy, &beta, iPforce->c, &incxy);

    //f'''=M^{1/2}*f''
    for( int i=0; i < _3N; i++) {
            iPforce->c[i] *= sqrtMass[i/3];
    }
    //put back into vector3DBlocks
    //vectTOvector3DBlock(tmpFX, iPforce);
    //delete temporary array
    return iPforce;
  }


  Vector3DBlock* NormalModeLangMF::slowspaceForce(Vector3DBlock * force, Vector3DBlock * iPforce){
    //
    if (force != iPforce)
      iPforce->intoAssign(*force);
    //vector3DBlockTOvect(force, tmpFX);        //put positions-x_0 into linear array
    //calculate M^{1/2}QQ^TM^{-1/2}f using BLAS
    //f'=M^{-1/2}*f
    for( int i=0; i < _3N; i++) {
            iPforce->c[i] *= invSqrtMass[i/3];
    }
    //c=Q^T*M^{-1/2}*f
    char transA = 'T';                                                  
// _rfM is the last rows of interest in Q
// _rfM-(firstFastMode-1) is the number of rows we need to define our subspace
// &((*Q)[_3N*(firstFastMode-1)]) (argument for degemv) lists the first row of interest

                                              // Transpose, LAPACK checks only first character N/V
    int m = _3N; int n = numSlowModes; int incxy = 1;     //sizes
    double alpha = 1.0; double beta = 0.0;
// FIXME: validate our args for dgemv
    Lapack::dgemv(&transA, &m, &n, &alpha, (*Q), &m, iPforce->c, &incxy, &beta, tmpC, &incxy);

    //f''=Qc
    char transB = 'N'; /* LAPACK checks only first character N/V */
    alpha = 1.0;        beta = 0.0;

    Lapack::dgemv(&transB, &m, &n, &alpha, (*Q), &m, tmpC, &incxy, &beta, iPforce->c, &incxy);

    //f'''=M^{1/2}*f''
    for( int i=0; i < _3N; i++) {
            iPforce->c[i] *= sqrtMass[i/3];
    }
    //put back into vector3DBlocks
    //vectTOvector3DBlock(tmpFX, iPforce);
    //delete temporary array
    return iPforce;
  }


  // FIXME: No validation has been performed on this function
  int NormalModeLangMF::minimizer(Real peLim, int numloop, bool simpM, bool reDiag,
    int *forceCalc, Real *lastLambda, ScalarStructure *myEnergies,
    Vector3DBlock *myPositions, GenericTopology *myTopo) 
  {
    int in,numLambda,itr,rsCG;
    Real oldPot,lastDiff,lambdaSlp1,a1,lambda,lambdaSlp,lambda1,calcPE;

    //initialize
    rsCG = 0; *lastLambda = 0.0; numLambda = 0; itr = 0; *forceCalc = 0; lastDiff = 5; 
    // start lastDiff at 5 kcal mol^{-1}
    // CHECK: rscg keeps track of ...? 

    utilityCalculateForces();
    (*forceCalc)++;
    //Set first value of /lambda to be 1/eigval
    lambda = 1.0 / *eigValP;    //exact solution for highest frequency mode if force mass weighted
    // lambda = distance we're moving the positions

    for(in=0;in<numloop;in++) {
      itr++;
      report.precision(10);
      report <<debug(6)<<"[NormalModeUtilities::minimizer] PE= "<<myEnergies->potentialEnergy()<<endr;
      // Here we *always* project onto the fast subspace
      fastspaceForce(myForces, myForces);
      // weight based on mass since we update positions directly
      for(int i=0;i<_N;i++) (*myForcesP)[i] /= myTopo->atoms[i].scaledMass; 
      //set posTemp
      posTemp = *myForcesP;

      //find slope of original PE with /lambda=0 here.
      lambdaSlp1 = 0.0;
      for( int k = 0; k < _N; k++ ) lambdaSlp1 -= posTemp[k].dot((*myForcesP)[k]);

      //save PE at /lambda=0
      oldPot = myEnergies->potentialEnergy();
      report <<debug(7)<<"[NormalModeUtilities::minimizer] lambd= "<<lambda<<endl;


      //find force at new position pos+/lambda*posTemp
      (*myPositions).intoWeightedAdd(lambda,posTemp);
         utilityCalculateForces();
      (*forceCalc)++;



      //Full minimizer? then solve for quadratic minimum

      if(!simpM){
        //find slope of PE with /lambda here
        lambdaSlp = 0.0;
        for( int k = 0; k < _N; k++ ) lambdaSlp -= posTemp[k].dot((*myForcesP)[k]);
        //solve for minimum for quadratic fit using two PE vales and the slope with /lambda
        Real a, b, oldLambda;
        oldLambda = lambda;
        a = -((myEnergies->potentialEnergy() - oldPot) / lambda - lambdaSlp) / lambda;
        b = lambdaSlp - 2.0 * a * lambda;
        lambda = -b / (2 * a);
        if(lambda <= 0.0) lambda = oldLambda;
        else{
          //Put solution into positions (but remove temporary solution for quadratic fit via oldLambda)
          (*myPositions).intoWeightedAdd(lambda-oldLambda,posTemp);
          utilityCalculateForces();
          (*forceCalc)++;
        }
      }

      //end full minimizer
      //update total gamma
      *lastLambda += lambda;
      numLambda++;

      //test for end, too large lambda test first
      if((oldPot - myEnergies->potentialEnergy()) < 0){
        // CHECK rsCG
        if(rsCG>4){
          report << error << "[NormalModeUtilities::minimizer] Minimization failed, Aborting. "<<rsCG <<endr;
        }
        else{
          if (!reDiag) {
	    a1 = (myEnergies->potentialEnergy() - oldPot - lambdaSlp1 * lambda) / (lambda * lambda);
      	    lambda1 = -lambdaSlp1 / (2 * a1);
	    // use the determinant
	    // jaw:  I'm hijacking from minimizer() in NormalModeUtilities and not validating the math
	    calcPE = a1*lambda1*lambda1+lambdaSlp1*lambda1+oldPot;

	    if(oldPot - calcPE > lastDiff && lambdaSlp1 != 0.0)
	      lambda1 = (-lastDiff * 2.0) / lambdaSlp1;

	    (*myPositions).intoWeightedAdd(-lambda,posTemp);
	    *lastLambda -= lambda;
	    numLambda--;
	    utilityCalculateForces();
	    (*forceCalc)++;
	    if(lambda1 > 0.0 && lambda1 < lambda) lambda = lambda1;
	    else lambda /= 2.0;
	    rsCG++;
	    report <<debug(1)<<"[NormalModeUtilities::minimizer] Reset CG, PE fail. Cycle= "<<rsCG<<" lambda= "<<lambda<<endl;
	  }
	  else {
	    (*myPositions).intoWeightedAdd(-lambda,posTemp);
	    report <<debug(1)<<"[NormalModeUtilities::minimizer] REDIAGONALIZING! PE fail  lambda= "<<lambda<<endl;
	    return -1;  // FIXME:  make sure we rediagonlize if we get this return
	  }

        }
      }
      else{
        rsCG = 0;
        lambda = 1.0 / *eigValP;    //revert to original value
      }
      if((oldPot - myEnergies->potentialEnergy()) < peLim && !rsCG) break;
      if(!rsCG) lastDiff = oldPot - myEnergies->potentialEnergy();
    }
    if(numLambda) *lastLambda /= (Real)numLambda;
    else lastLambda = 0;
    return itr;

  }

  // Project forces onto fast subspace, average the slow forces
  void NormalModeLangMF::forceProjection(bool bClear){
    unsigned int count = myForces->size();
    if (bClear) aveForceCount = 0;
    if((*Q) != NULL){
      // myForces has total forces
      slowForces.resize(count);
      slowForces.intoAssign(*myForces);
      // project myForces onto fast subspace
      fastspaceForce(myForces, myForces);
      // difference between old myForces stored in slowForces and myForces
      // gives us instantaneous slow force
      for( unsigned int i=0;i < count; i++) slowForces[i] -= (*myForces)[i];
      // add slow force to running sum in tempV3DBlk[i]
      if (bClear) slowForcesAvg.intoAssign(*slowForces);
      else {
	for( unsigned int i=0;i < count; i++) slowForcesAvg[i] += slowForces[i];
      }
      aveForceCount++;
    }
    //slowForcesAvg/aveForcecout ??
  }
  void NormalModeLangMF::forceAvg() {
    for( unsigned int i=0;i < count; i++) slowForcesAvg[i] /= aveForceCount;
  }

  void NormalModeLangMF::genFastGauss(Vector3DBlock *gaussRandCoord, GenericTopology *myTopo) {
    //generate set of random force variables and project into sub space
    for( int i = 0; i < _3N; i++ )
        (*gaussRandCoord)[i/3][i%3] = randomGaussianNumber(mySeed);//
    for( int i = 0; i < _N; i++ )
        (*gaussRandCoord)[i] *= sqrtMass[i];
    fastsubspaceForce(gaussRandCoord, gaussRandCoord);
    for( int i = 0; i < _N; i++ )
        (*gaussRandCoord)[i] /= myTopo->atoms[i].scaledMass;
  }

  void NormalModeUtilities::genSlowGauss(Vector3DBlock *gaussRandCoord, GenericTopology *myTopo) {
    //generate set of random force variables and project into sub space
    for( int i = 0; i < _3N; i++ )
        (*gaussRandCoord)[i/3][i%3] = randomGaussianNumber(mySeed);//
    for( int i = 0; i < _N; i++ )
        (*gaussRandCoord)[i] *= sqrtMass[i];
    slowspaceForce(gaussRandCoord, gaussRandCoord);
    for( int i = 0; i < _N; i++ )
        (*gaussRandCoord)[i] /= myTopo->atoms[i].scaledMass;
  }


  // FIXME:  need a set of validation tests for this function
  void NormalModeLangMF::run(int numTimesteps,int numFastTimesteps)
  {
    int i,j,itrs,forceCalc;
    Real lastLambda;
    forceCalc = 0;

    //check valid eigenvectors
    if(*Q == NULL)
        report << error << "No Eigenvectors for NormalMode integrator."<<endr;

    // 1.  propagate over the slow subspace
// FIXME: make sure numTimsteps and numFastTimesteps
// are at the correct values
    for(i=0;i<numTimesteps;i++) {

      //preStepModify();
      if (i > 0) {
        genSlowGauss(&gaussRandCoord1, app->topology);
        doHalfKickSlow();
        doDrift();
// FIXME: where do the gaussian projections go?
//        genProjGauss(&gaussRandCoord1, app->topology);
      }
      //postStepModify();

      // 2. propagate over the fast subspace
      if (minimizeFastSpace) { 
        itrs = minimizer(minLim, maxMinSteps, true, false, &forceCalc, &lastLambda,
                &app->energies, &app->positions, app->topology);
      }


      for(j=0;j<numFastTimesteps;j++) {

        //preStepModify();
	if (j == 0) {
          calculateForces();
          // project onto the fast subspace, average the slow forces
          forceProjection(TRUE);

	}	
 	genFastGauss(&gaussRandCoord1, app->topology);
        doHalfKickFast();
        doDrift();

        // calculate forces
        calculateForces(FALSE);
        // project onto the fast subspace, average the slow forces
        forceProjection();

// FIXME:  redo function for Gaussian noise - fast/slow subspaces
// FIXME:  validation test:  check convergence of MF with and without Gaussian noise
        genFastGauss(&gaussRandCoord1, app->topology);
        doHalfKickFast();
        //postStepModify();
        

      }
      forceAvg();
      if (i > 0) {
        genSlowGauss(&gaussRandCoord1, app->topology);
        doHalfKickSlow();
      }
    }
  }

  void NormalModeLangLf::doDrift() {
      const Real h = getTimestep() * Constant::INV_TIMEFACTOR;
      app->positions.intoWeightedAdd(h, app->velocities);
      buildMolecularCenterOfMass(&app->positions, app->topology);
  }

  void NormalModeLangLf::doHalfKickSlow() {
    const Real dt = getTimestep() * Constant::INV_TIMEFACTOR; // in fs
    const Real fdt = ( 1.0 - exp( -0.5 * myGamma * dt ) ) / myGamma;
    const Real vdt = exp(-0.5*myGamma*dt);
       ohys.velocties = phys.velocities*numpy.sqrt(phys.invmasses)
    const Real ndt = sqrt( ( 1.0 - exp( -myGamma * dt ) ) / (2.0 * myGamma) ); //was sqrt( fdt );
    const Real sqrtFCoverM = sqrt( 2.0 * Constant::BOLTZMANN * myTemp * myGamma );

    for( int i = 0; i < _N; i++ ) {
        // semi-update velocities
        app->velocities[i] = app->velocities[i]*vdt
                                +(*myForces)[i] * fdt / app->topology->atoms[i].scaledMass
                                    +gaussRandCoord1[i]*sqrtFCoverM*ndt;
    }
    nonSubspacePosition(&app->velocities, &app->velocities);
    buildMolecularMomentum(&app->velocities, app->topology);
  }

  void NormalModeLangLf::doHalfKickFast() {
    const Real dt = getTimestep() * Constant::INV_TIMEFACTOR; // in fs
    const Real fdt = ( 1.0 - exp( -0.5 * myGamma * dt ) ) / myGamma;
    const Real vdt = exp(-0.5*myGamma*dt);
       ohys.velocties = phys.velocities*numpy.sqrt(phys.invmasses)
    const Real ndt = sqrt( ( 1.0 - exp( -myGamma * dt ) ) / (2.0 * myGamma) ); //was sqrt( fdt );
    const Real sqrtFCoverM = sqrt( 2.0 * Constant::BOLTZMANN * myTemp * myGamma );

    for( int i = 0; i < _N; i++ ) {
        // semi-update velocities
        app->velocities[i] = app->velocities[i]*vdt
                                +(*myForces)[i] * fdt / app->topology->atoms[i].scaledMass
                                    +gaussRandCoord1[i]*sqrtFCoverM*ndt;
    }
    subspaceVelocity(&app->velocities, &app->velocities);
    buildMolecularMomentum(&app->velocities, app->topology);
  }

//change the parameters
  void NormalModeLangMF::getParameters(vector<Parameter>& parameters) const {
    STSIntegrator::getParameters(parameters);
    parameters.push_back(Parameter("firstMode",Value(firstMode,ConstraintValueType::NotNegative()),1,Text("First mode of the slow subspace")));
    parameters.push_back(Parameter("numbermodes",Value(numMode,ConstraintValueType::NotNegative()),1,Text("Number of modes propagated")));
    parameters.push_back(Parameter("gamma",Value(myGamma*(1000 * Constant::INV_TIMEFACTOR),ConstraintValueType::NotNegative()),80.0,Text("Langevin Gamma")));
    parameters.push_back(Parameter("seed",Value(mySeed,ConstraintValueType::NotNegative()),1234,Text("Langevin random seed")));
    parameters.push_back(Parameter("temperature",Value(myTemp,ConstraintValueType::NotNegative()),300.0,Text("Langevin temperature")));
    parameters.push_back(Parameter("gencompnoise",Value(genCompNoise,ConstraintValueType::NoConstraints()),false,Text("Generate complimentary noise")));
    parameters.push_back(Parameter("minimizeFastSpace",Value(minimizeFastSpace,ConstraintValueType::NoConstraints()),true,Text("Minimize the fast subspace")));
    parameters.push_back(Parameter("numFastTimesteps",Value(numFastTimesteps,ConstraintValueType::NotNegative()),10,Text("Number of iterations for the fast subspace")));
 }

//change the values
  STSIntegrator* NormalModeLangMF::doMake(const vector<Value>& values,ForceGroup* fg)const{
    return new NormalModeLangMF(values[0],values[1],values[2],values[3],values[4],values[5],values[6],fg);
  }


}
